'''
author:        Wang Chenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2024-11-20
Copyright © Department of Physics, Tsinghua University. All rights reserved

'''

import bibtexparser
import pickle
import json

def find_info_from_bibfile(file_path):
    ''' Find title, abstract and keywords from bibfile
        if an item has no abstract or keywords, it will be ignored
    '''

    with open(file_path, 'r', encoding='utf-8') as bibtex_file:
        bib_database = bibtexparser.load(bibtex_file)
    
    all_info = []
    for entry in bib_database.entries:
        title = entry.get("title", '')
        abstract = entry.get("abstract", '')
        keywords = entry.get("keywords", '')
        if title == '' or abstract == '' or keywords == '':
            continue
        all_info.append({
            "title": title,
            "abstract": abstract,
            "keywords": keywords
        })
    
    print(len(all_info))
    
    with open('data/info.pkl', 'wb') as f:
        pickle.dump(all_info, f)


def generate_prompt(all_info, output_file_path):
    ''' Generate prompt for LLM'''
    system_prompt = (
        "[Character settings]You are a physicist and English native speaker."
        " Now you need to read the title and the abstract of an article, and"
        " extract keywords for each literature. You need to search from \"\"\"{{knowledge}}\"\"\" to find better keywords.\n\n\n "
        "[order]\n1. Keywords are seperated by ','.\n"
        "2.Please search the knowledge base. If there are suitable keywords in the knowledge base, "
            "return the keywords in the knowledge base first.\n"
        "3.If one keyword <kw1> has a parent keyword <kw2>, please return in the format of \"<kw2>/<kw1>\"\n"
        "4.**GIVE ME KEYWORDS ONLY, DO NOT REPEAT TITLE AND ABSTRACT**\n"
        "5.Find keywords from \n\"\"\"\n{{knowledge}}\n\"\"\"\n. If there's no suitable keywords, please conclude from the title and abstract, and the keywords must end with '(ai)'\n\n"
    )

    with open("LLM-config.json", "r") as fp:
        LLM_json = json.load(fp)
        user_key = LLM_json['LLM-key']
        user_chat_model = LLM_json['LLM-name']
        knowledge_base_id = LLM_json['KB-id']

    with open(output_file_path, 'w', encoding='utf-8') as f:
        pass
    for item in all_info:
        msg = [
            {"role": "user", "content": (
                "[Title] %s, " 
                "[Abstract] %s" % (item["title"], item["abstract"])
            )},
            {"role": "assistant", "content": item["keywords"]},
        ]
        tools=[
            {
                "type": "retrieval",
                "retrieval": {
                    "knowledge_id": knowledge_base_id,
                    "prompt_template": system_prompt
                }
            }
        ]
        json_data = {"messages": msg, "tools": tools}
        with open(output_file_path, 'a', encoding='utf-8') as f:
            json.dump(json_data, f)
            f.write('\n')


def main():
    find_info_from_bibfile('data/articles-with-abstract.bib')
    with open('data/info.pkl', 'rb') as f:
        all_info = pickle.load(f)
    generate_prompt(all_info, 'data/LLM-data.jsonl')


if __name__ == "__main__":
    main()